#include <stdio.h>
#include <string.h>
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "esp_netif.h"
#include "esp_heap_caps.h"
#include "ESP32_wifi_connect.h"

static const char *TAG = "WiFi连接模块";
static bool wifi_connected = false;
static int wifi_retry_count = 0;
static esp_netif_t *sta_netif = NULL;
#define MAX_WIFI_RETRY 5

// 安全清理函数
static void wifi_cleanup() {
    if (wifi_connected) {
        esp_wifi_disconnect();
    }
    esp_wifi_stop();
    esp_wifi_deinit();
    
    if (sta_netif) {
        esp_netif_destroy(sta_netif);
        sta_netif = NULL;
    }
    
    wifi_connected = false;
}

// 事件处理函数
void wifi_event_handler(void* arg, esp_event_base_t event_base, 
                      int32_t event_id, void* event_data)
{
    if (event_base == WIFI_EVENT) {
        switch (event_id) {
            case WIFI_EVENT_STA_START:
                ESP_LOGI(TAG, "WiFi驱动就绪，开始连接...");
                esp_wifi_connect();
                break;
                
            case WIFI_EVENT_STA_CONNECTED:
                ESP_LOGI(TAG, "已关联到AP");
                break;
                
            case WIFI_EVENT_STA_DISCONNECTED: {
                wifi_event_sta_disconnected_t* disconn = (wifi_event_sta_disconnected_t*) event_data;
                ESP_LOGW(TAG, "连接断开，原因: %d - %s", 
                    disconn->reason,
                    disconn->reason == WIFI_REASON_AUTH_EXPIRE ? "认证超时" :
                    disconn->reason == WIFI_REASON_ASSOC_LEAVE ? "主动断开" : "其他错误");
                
                if (wifi_retry_count < MAX_WIFI_RETRY) {
                    vTaskDelay(pdMS_TO_TICKS(1000));
                    esp_wifi_connect();
                    wifi_retry_count++;
                    ESP_LOGI(TAG, "第%d次重试连接...", wifi_retry_count);
                } else {
                    ESP_LOGE(TAG, "超过最大重试次数(%d)", MAX_WIFI_RETRY);
                }
                wifi_connected = false;
                break;
            }
            default:
                break;
        }
    } else if (event_base == IP_EVENT) {
        switch (event_id) {
            case IP_EVENT_STA_GOT_IP: {
                ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
                ESP_LOGI(TAG, "获取IP成功: " IPSTR "/" IPSTR " 网关: " IPSTR,
                    IP2STR(&event->ip_info.ip),
                    IP2STR(&event->ip_info.netmask),
                    IP2STR(&event->ip_info.gw));
                wifi_connected = true;
                wifi_retry_count = 0;
                break;
            }
            case IP_EVENT_STA_LOST_IP:
                ESP_LOGW(TAG, "IP地址丢失");
                wifi_connected = false;
                break;
        }
    }
}

// WiFi初始化函数
esp_err_t wifi_init_sta(const char *ssid, const char *password)
{
    if (!ssid || !password || !*ssid || !*password) {
        ESP_LOGE(TAG, "SSID或密码无效");
        return ESP_ERR_INVALID_ARG;
    }

    // 清理旧配置
    wifi_cleanup();

    // 初始化网络栈
    static bool netif_inited = false;
    if (!netif_inited) {
        ESP_ERROR_CHECK(esp_netif_init());
        ESP_ERROR_CHECK(esp_event_loop_create_default());
        netif_inited = true;
    }

    // 创建网络接口
    sta_netif = esp_netif_create_default_wifi_sta();
    if (!sta_netif) {
        ESP_LOGE(TAG, "创建网络接口失败");
        return ESP_FAIL;
    }

    // WiFi初始化配置
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    cfg.nvs_enable = false; // 禁用NVS存储配置
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // 注册事件处理器
    ESP_ERROR_CHECK(esp_event_handler_instance_register(
        WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL, NULL));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(
        IP_EVENT, IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL, NULL));

    // 配置WiFi参数
    wifi_config_t wifi_config = {
        .sta = {
            .threshold.authmode = WIFI_AUTH_WPA2_PSK,
            .pmf_cfg = {.capable = true, .required = false},
        },
    };
    strlcpy((char*)wifi_config.sta.ssid, ssid, sizeof(wifi_config.sta.ssid));
    strlcpy((char*)wifi_config.sta.password, password, sizeof(wifi_config.sta.password));

    // 启动WiFi
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG, "启动WiFi连接，目标AP: %s", ssid);
    return ESP_OK;
}

// 连接状态检查
bool wifi_is_connected(void) {
    return wifi_connected;
}

// 主动断开连接
void wifi_disconnect(void) {
    wifi_cleanup();
    ESP_LOGI(TAG, "已主动断开WiFi连接");
}
